%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%
%%%%	Navier-Stokes problem
%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Author: Alonso Alvarez

%BEGIN PROBLEM DEFINITION
%-------------------------------------------------------------
%Convention used
% Ny   x  x  x  ... x
% .    .
% .    .
% .    .
% 3    x  x  x  ... x
% 2    x  x  x  ... x
% 1    x  x  x  ... x
%
% j/i  1  2  3  ... Nx

global Nx
Nx = 240;   %X Direction
global Ny
Ny = 20;    %Y Direction
Lx = 12.0;
Ly = 1.0;
visc = 0.1;
h = Ly/Ny
rho = 1.0;
sor = 0.8;
% initial data
u = zeros((Nx-1)*Ny,1);
v = zeros(Nx*(Ny-1),1);
p = zeros(Nx*Ny, 1);
%
unovo = zeros((Nx-1)*Ny,1);
vnovo = zeros(Nx*(Ny-1),1);
%
%
utemp = zeros((Nx-1)*Ny,1);
vtemp = zeros(Nx*(Ny-1),1);
%

for i = 1: Nx-1
    for j = 1: Ny
        %u(uij(i,j),1) = 4.*j*h*(1 - j*h);
    end
end

%temporal integration
dt = 0.01;
T0 = 0.;
T1 = 1.;

NumTime = (T1 - T0)/dt;

%% Matrix for pressure
Ap = zeros(Nx*Ny, Nx*Ny);
bp = zeros(Nx*Ny, 1);

for i = 1: Nx
    for j = 1: Ny
        centro = 4;
        if( j > 1)  Ap(pij(i,j),pij(i,j-1)) = 1.0; else centro = centro - 1; end
        if( j < Ny) Ap(pij(i,j),pij(i,j+1)) = 1.0; else centro = centro - 1; end
        if( i > 1)  Ap(pij(i,j),pij(i-1,j)) = 1.0; else centro = centro - 1; end
        if( i < Nx) Ap(pij(i,j),pij(i+1,j)) = 1.0; else centro = centro + 1; end
        Ap(pij(i,j),pij(i,j)) = -centro;
    end
end
Ap = Ap/(h^2);
Ap = sparse(Ap);

%%%% Temporal loop
for tt = 1: NumTime
    %print('     SOLVING.. %d iteration\n', tt);
    
    %%%%   Predictor step   %%%%%%%
    for i = 1: Nx
        for j = 1: Ny
            
            if(i < Nx && j < Ny)
                upto = u(uij(i,j),1);
                vpto = v(vij(i,j),1);
                
                if(i > 1)
                    ule = u(uij(i-1,j),1);
                    vle = v(vij(i-1,j),1);
                    udi = u(uij(i-1,j+1),1);
                else
                    ule = 4*j*h*(1 - j*h);
                    udi = 4*(j+1)*h*(1 - (j+1)*h);
                    vle = -vpto;
                end
                
                if(i < Nx - 1) uri = u(uij(i+1,j),1); else uri = upto; end
                vri = v(vij(i+1,j),1);
                
                if( j > 1 )
                    udo = u(uij(i,j-1),1);
                    vdo = v(vij(i,j-1),1);
                    vdi = v(vij(i+1,j-1),1);
                else
                    udo = -upto;
                    vdo = 0.0;
                    vdi = 0.0;
                end
                
                if( j < Ny - 1) vup = v(vij(i,j+1),1); 	else vup = 0.0; end
                uup = u(uij(i,j+1),1);
            end
            
            if (i == Nx && j < Ny)
                vpto = v(vij(i,j),1);
                vle = v(vij(i-1,j),1);
                vri = -vpto;
                if( j > 1 ) vdo = v(vij(i,j-1),1); 	else vdo = 0.0; end
                if( j < Ny - 1 ) vup = v(vij(i,j+1),1); else vup = 0.0; end
                vdi = -vdo;
            end
            
            if ( j == Ny && i < Nx )
                upto = u(uij(i,j),1);
                udo = u(uij(i,j-1),1);
                uup = -upto;
                
                % Parábola
                %if( i > 1 ) ule = u(uij(i-1,j),1); else ule = 4*j*h*(1 - j*h); end
                % Frente Plano
                if( i > 1 ) ule = u(uij(i-1,j),1); else ule = 1; end
                
                if( i < Nx - 1) uri = u(uij(i+1,j),1); 	else uri = -upto; end
                udi = -ule;
            end
            
            Ax = ((uri + upto)^2 - (upto + ule)^2 + (uup + upto)*(vri + vpto) - (upto + udo)*(vdi + vdo))/(4*h);
            Ay = ((vup + vpto)^2 - (vpto + vdo)^2 + (uup + upto)*(vri + vpto) - (udi + ule)*(vpto + vle))/(4*h);
            
            Dx = (uri + ule + uup + udo - 4*upto)/(h^2);
            Dy = (vri + vle + vup + vdo - 4*vpto)/(h^2);
            
            if(i < Nx)
                unovo(uij(i,j),1) = upto + dt*(-Ax + visc*Dx);
            end
            if(j < Ny)
                vnovo(vij(i,j),1) = vpto + dt*(-Ay + visc*Dy);
            end
        end
    end
    
    %%%%   Pressure Poisson equation  %%%%%%%
    for i = 1: Nx
        for j = 1: Ny
            
            if(i < Nx) uri = unovo(uij(i,j)); else uri = unovo(uij(i-1,j)); end
            % Parábola
            %if(i > 1 ) ule = unovo(uij(i-1,j)); else ule = 4*j*h*(1 - j*h); end
            % Frente Plano
            if(i > 1 ) ule = unovo(uij(i-1,j)); else ule = 1; end
            if(j < Ny) vup = vnovo(vij(i,j)); else vup = 0.; end
            if(j > 1 ) vdo = vnovo(vij(i,j-1)); else vdo = 0.; end
            
            bp(pij(i,j),1) = rho*(uri - ule + vup - vdo)/(dt*h);
        end
    end
    p = Ap\bp;
    
    %%%%   Predictor step  %%%%%%%
    for i = 1: Nx
        for j = 1: Ny
            ppto = p(pij(i,j));
            if(i < Nx)
                upto = unovo(uij(i,j),1);
                pri = p(pij(i+1,j),1);
                utemp(uij(i,j),1) = upto - dt*(pri - ppto)/(rho*h);
            end
            if(j < Ny)
                vpto = vnovo(vij(i,j),1);
                pup = p(pij(i,j+1),1);
                vtemp(vij(i,j),1) = vpto - dt*(pup - ppto)/(rho*h);
            end
        end
    end
    u = sor*utemp + (1 - sor)*u;
    v = sor*vtemp + (1 - sor)*v;
end


%%%%%%%%%%%%%%   BEGIN POST-PROCESSING STEP

%% Horizontal velocity
X = zeros(Nx-1,1);
for i = 1 : Nx - 1
    X(i) = h*i;
end
Y = zeros(Ny,1);
Y(1) = h/2;
for j = 2: Ny
    Y(j) = Y(j-1) + h;
end

figure()
U = reshape(u(:,1),Ny,Nx-1);
surf(X,Y,U);
%axis([0 2 0 1 0 1]);

%% Vertical velocity
X = zeros(Nx,1);
X(1) = h/2;
for i = 2 : Nx
    X(i) = X(i-1) + h;
end
Y = zeros(Ny-1,1);
for j = 1: Ny-1
    Y(j) = j*h;
end

figure()
U = reshape(v(:,1),Ny-1,Nx);
surf(X,Y,U);
%axis([0 2 0 1 0 1]);
title('Horizontal')

%% Pressure
X = zeros(Nx,1);
X(1) = h/2;
for i = 2 : Nx
    X(i) = X(i-1) + h;
end
Y = zeros(Ny,1);
Y(1) = h/2;
for j = 2: Ny
    Y(j) = Y(j-1) + h;
end
title('Vertical')

figure()
U = reshape(p(:,1),Ny,Nx);
surf(X,Y,U);
%axis([0 2 0 1 0 1]);
pause();
